feat(megatron): report peak memory for log_probs and actor_train - #2304
Open
EazyReal wants to merge 1 commit into
Open
feat(megatron): report peak memory for log_probs and actor_train#2304EazyReal wants to merge 1 commit into
EazyReal wants to merge 1 commit into
Conversation
print_memory samples current allocation only at fixed lifecycle points, so it cannot show a phase high-water mark and an OOM can happen before the next report. Add accelerator-aware peak reporting in a finally block and compose it with the existing log_probs and actor_train timers. Accelerators without peak-memory support keep executing without the report. CUDA/ROCm and MUSA expose the required allocator APIs. Measured on a 128x H100 DeepSeek-V4-Flash run at 262,144-token context: the finally-path report attributed a 132k-context log_probs OOM to 25 GiB of retained activations.
EazyReal
marked this pull request as ready for review
August 22, 2026 21:41
4 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
print_memorysamples current allocation only at fixed lifecycle points, so it cannot expose a phase's high-water mark and an OOM can happen before the next reportreport_peak_memory(phase): reset peak stats on entry and log peak allocated/reserved memory in afinallyblock, so an OOMing phase still reports its high-water marklog_probsforward andactor_train, composed with the existing timer scopes; a reused or skipped actor forward emits no separatelog_probsreportMeasured on a 128x H100 DeepSeek-V4-Flash run at 262,144-token context: the
finally-path report attributed a 132k-contextlog_probsOOM to 25 GiB of retained activations. The OOM traceback alone did not show the phase's high-water mark.Testing
python tests/test_accelerator.py— 12 passed; covers reset-before-body, exact GB/log formatting, reporting on exceptions without swallowing them, unsupported-accelerator fallback, and CUDA/MUSA capability detectionblack --line-length 119,ruff check, andisort --profile=blackpass on all touched files